package com.bff.gaia.optimizer.joinarrangement;

import com.bff.gaia.api.common.ExecutionConfig;
import com.bff.gaia.api.common.ExecutionMode;
import com.bff.gaia.api.common.functions.ProjectFlatJoinFunction;
import com.bff.gaia.api.common.operators.BinaryOperatorInformation;
import com.bff.gaia.api.common.operators.base.InnerJoinTreeOperatorBase;
import com.bff.gaia.api.common.operators.base.JoinOperatorBase;
import com.bff.gaia.api.common.typeinfo.TypeInformation;
import com.bff.gaia.api.common.typeinfo.Types;
import com.bff.gaia.api.java.typeutils.TupleTypeInfo;
import com.bff.gaia.optimizer.dag.DataSinkNode;
import com.bff.gaia.optimizer.dag.JoinTreeNode;
import com.bff.gaia.optimizer.dag.OptimizerNode;
import com.bff.gaia.optimizer.traversals.IdAndEstimatesVisitor;

import java.util.*;

public class JoinSequenceArrangement {

	private final int defaultParallelism;
	private final ExecutionMode defaultDataExchangeMode;
	private final ExecutionConfig config;

	private IdAndEstimatesVisitor.JoinTreeRecord joinTreeRecord;
	private Map<Integer, OptimizerNode> joinTreeNodeMap;
	private Map<Integer, OptimizerNode> joinTreeSourceNodeMap;

	public JoinSequenceArrangement(int defaultParallelism, ExecutionMode defaultDataExchangeMode, ExecutionConfig config, IdAndEstimatesVisitor.JoinTreeRecord joinTreeRecord) {

		this.defaultParallelism = defaultParallelism;
		this.defaultDataExchangeMode = defaultDataExchangeMode;
		this.config = config;

		this.joinTreeRecord = joinTreeRecord;
		this.joinTreeNodeMap = joinTreeRecord.getJoinTreeNodeMap();
		this.joinTreeSourceNodeMap = joinTreeRecord.getJoinTreeSourceNodeMap();

	}

	public void arrangeJoinNodeSequence() {

		joinTreeRecord.initKeySetList();
		for (IdAndEstimatesVisitor.JoinTreeRecord.ColumnSourceNodeSet keySet : joinTreeRecord.getKeySetList()) {
			for (IdAndEstimatesVisitor.JoinTreeRecord.ColumnSourceNode keyNode : keySet.getColumnSourceNodeSet()) {
				System.out.println("nodeId:" + keyNode.getNodeId() + "\t" + keyNode.getColumnPos() + "\tsubId:" + keyNode.getSuperId());
			}
			System.out.println("===============================================");
		}
		System.out.println("initKeySetList() complete\n");
		joinTreeRecord.initResultSetList();
		for (IdAndEstimatesVisitor.JoinTreeRecord.ColumnSourceNodeSet resultSet : joinTreeRecord.getResultSetList()) {
			for (IdAndEstimatesVisitor.JoinTreeRecord.ColumnSourceNode resultNode : resultSet.getColumnSourceNodeSet()) {
				System.out.println("pos:" + joinTreeRecord.getResultSetList().indexOf(resultSet) + "\tnodeId:" + resultNode.getNodeId() + "\t" + resultNode.getColumnPos() + "\tsubId:" + resultNode.getSuperId());
			}
			System.out.println("===============================================");
		}
		System.out.println("initResultSetList() complete\n");

		DynamicPlan dynamicPlan = new DynamicPlan(joinTreeRecord);
		DynamicPlan.Tree root = dynamicPlan.reArrange();
		traverseTree(root,true);

	}

	@SuppressWarnings("unchecked")
	private JoinTreeNode restructureJoinTreeNode(OptimizerNode input1, OptimizerNode input2, boolean isTailNode) {

		int keyPos1;
		int keyPos2 = 0;
		int[] keySetPos1 = getColumnInfo(input1).getKeySetPos();
		int[] keySetPos2 = getColumnInfo(input2).getKeySetPos();

		//生成keyPos1 keyPos2
		search:
		for (keyPos1 = 0; keyPos1 < keySetPos1.length; keyPos1++) {
			for (keyPos2 = 0; keyPos2 < keySetPos2.length; keyPos2++) {
				if (keySetPos1[keyPos1] == keySetPos2[keyPos2] && keySetPos1[keyPos1] != -1) {
					break search;
				}
			}
		}

		if (keyPos1 == keySetPos1.length || keySetPos1[keyPos1] != keySetPos2[keyPos2]) {
			throw new IllegalArgumentException(input1 + "cannot join" + input2);
		}

		//将已加入的node标为finished
		Set<Integer> finishedNodeSet = new HashSet<>();
		finishedNodeSet.addAll(getColumnInfo(input1).getFinishedNodeSet());
		finishedNodeSet.addAll(getColumnInfo(input2).getFinishedNodeSet());

		for (Integer nodeId : finishedNodeSet) {
			joinTreeRecord.setNodeFinished(joinTreeRecord.getKeySetList(), nodeId);
		}

		//生成新的ColumnInfo
		int[] resultSetPos1 = getColumnInfo(input1).getResultPos();
		int[] resultSetPos2 = getColumnInfo(input2).getResultPos();

		ArrayList<Integer> reserveColumn1 = new ArrayList<>();
		ArrayList<Integer> reserveColumn2 = new ArrayList<>();
		ArrayList<Integer> newKeySetPos = new ArrayList<>();
		ArrayList<Integer> newResultSetPos = new ArrayList<>();

		for (int i = 0; i < keySetPos1.length; i++) {
			if (resultSetPos1[i] != -1) {
				reserveColumn1.add(i);
				newKeySetPos.add(keySetPos1[i]);
				newResultSetPos.add(resultSetPos1[i]);
			} else if (keySetPos1[i] != -1) {
				if (!joinTreeRecord.getKeySetList().get(keySetPos1[i]).checkFinished()) {
					reserveColumn1.add(i);
					newKeySetPos.add(keySetPos1[i]);
					newResultSetPos.add(-1);
				}
			}
		}

		for (int i = 0; i < keySetPos2.length; i++) {
			if (resultSetPos2[i] != -1) {
				if (i != keyPos2) {
					reserveColumn2.add(i);
					newKeySetPos.add(keySetPos2[i]);
					newResultSetPos.add(resultSetPos2[i]);
				}
			} else if (keySetPos2[i] != -1) {
				if (!joinTreeRecord.getKeySetList().get(keySetPos2[i]).checkFinished()) {
					reserveColumn2.add(i);
					newKeySetPos.add(keySetPos2[i]);
					newResultSetPos.add(-1);
				}
			}
		}

		InnerJoinTreeOperatorBase.ColumnInfo columnInfo = new InnerJoinTreeOperatorBase.ColumnInfo(newKeySetPos, newResultSetPos, finishedNodeSet);

		//生成fields和isFromFirst
		int[] fields = new int[reserveColumn1.size() + reserveColumn2.size()];
		boolean[] isFromFirst = new boolean[reserveColumn1.size() + reserveColumn2.size()];
		if (isTailNode) {
			for (int column1 : reserveColumn1) {
				isFromFirst[resultSetPos1[column1]] = true;
				fields[resultSetPos1[column1]] = column1;
			}
			for (int column2 : reserveColumn2) {
				fields[resultSetPos2[column2]] = column2;
			}
		} else {
			ArrayList<Integer> column = new ArrayList<>();
			column.addAll(reserveColumn1);
			column.addAll(reserveColumn2);
			fields = column.stream().mapToInt(Integer::intValue).toArray();
			Arrays.fill(isFromFirst, 0, reserveColumn1.size(), true);
		}

		//生成OperatorInformation
		TypeInformation inputType1 = input1.getOperator().getOperatorInfo().getOutputType();
		TypeInformation inputType2 = input2.getOperator().getOperatorInfo().getOutputType();

		ArrayList<TypeInformation> outputTypeList = new ArrayList<>();
		for (int i = 0; i < fields.length; i++) {
			outputTypeList.add(((TupleTypeInfo) (isFromFirst[i] ? inputType1 : inputType2)).getTypeAt(fields[i]));
		}
		TypeInformation outputType = Types.TUPLE((TypeInformation[]) outputTypeList.toArray(new TypeInformation[outputTypeList.size()]));

		BinaryOperatorInformation operatorInfo = new BinaryOperatorInformation(inputType1, inputType2, outputType);

		//生成udf
		ProjectFlatJoinFunction udf = new ProjectFlatJoinFunction(fields, isFromFirst, ((TupleTypeInfo) outputType).createSerializer(config).createInstance());

		//生成InnerJoinTreeOperatorBase
		InnerJoinTreeOperatorBase innerJoinTreeOperatorBase = new InnerJoinTreeOperatorBase(udf, operatorInfo, new int[]{keyPos1}, new int[]{keyPos2}, "joinTree");

		innerJoinTreeOperatorBase.setFirstInput(input1.getOperator());
		innerJoinTreeOperatorBase.setSecondInput(input2.getOperator());
		innerJoinTreeOperatorBase.setParallelism(-1);
		innerJoinTreeOperatorBase.setCustomPartitioner(null);
		innerJoinTreeOperatorBase.setJoinHint(JoinOperatorBase.JoinHint.OPTIMIZER_CHOOSES);

		innerJoinTreeOperatorBase.setColumnInfo(columnInfo);
		innerJoinTreeOperatorBase.setSemanticProperties(isFromFirst, fields);

		//生成JoinTreeNode
		JoinTreeNode joinTreeNode = new JoinTreeNode(innerJoinTreeOperatorBase);

		joinTreeNode.setParallelism(this.defaultParallelism);
		joinTreeNode.setInput(input1, input2, this.defaultDataExchangeMode);

		return joinTreeNode;
	}

	private InnerJoinTreeOperatorBase.ColumnInfo getColumnInfo(OptimizerNode inputNode) {
		if (inputNode instanceof JoinTreeNode) {
			return ((InnerJoinTreeOperatorBase) inputNode.getOperator()).getColumnInfo();
		} else {
			int joinTreeSourceNodeId = inputNode.getJoinTreeSourceNodeId();
			int columnLength = ((TupleTypeInfo) inputNode.getOperator().getOperatorInfo().getOutputType()).getTypeLength();
			int[] keySetPos = new int[columnLength];
			int[] resultSetPos = new int[columnLength];
			for (int i = 0; i < columnLength; i++) {
				keySetPos[i] = joinTreeRecord.getSuperSetId(joinTreeRecord.getKeySetList(), joinTreeSourceNodeId, i);
				resultSetPos[i] = joinTreeRecord.getSuperSetId(joinTreeRecord.getResultSetList(), joinTreeSourceNodeId, i);
			}
			return new InnerJoinTreeOperatorBase.ColumnInfo(keySetPos, resultSetPos, joinTreeSourceNodeId);
		}
	}

	private OptimizerNode traverseTree(DynamicPlan.Tree tree, boolean isTailNode) {
		DynamicPlan.Tree child1 = tree.child1;
		DynamicPlan.Tree child2 = tree.child2;

		if (child1 == null && child2 == null) {
			return joinTreeSourceNodeMap.get(tree.getId());
		}

		JoinTreeNode joinTreeNode = restructureJoinTreeNode(traverseTree(child1, false), traverseTree(child2, false), isTailNode);
		joinTreeNode.setEstimatedNumRecords(tree.getEstimatedNumRecords());
		joinTreeNode.setEstimatedOutputSize(tree.getEstimatedOutputSize());

		if (isTailNode) {
			OptimizerNode afterNode = (joinTreeNodeMap.get(joinTreeNodeMap.size() - 1).getOutgoingConnections().get(0)).getTarget();
			if (afterNode instanceof DataSinkNode) {
				((DataSinkNode) afterNode).reSetInput(joinTreeNode, this.defaultDataExchangeMode);
			}
		}

		System.out.println(tree);

		return joinTreeNode;
	}

}